import os
import time

import pandas as pd

from Config import *

pd.set_option('expand_frame_repr', False)  # 当列太多时不换行


# ===时间格式转换函数===
def time_strptime_timestamp(time_para, time_format):
    # time_format = '%Y-%m-%d %H:%M:%S'
    # 将UNIX时间戳转化为标准时间(UTC)；入参time_type如：1609430400000(ms级别)
    dt = datetime.fromtimestamp(time_para / 1000)  # 将时间戳转换为datetime对象
    utc_dt = dt.astimezone(timezone.utc)  # 将时间转换为UTC时间
    res = utc_dt.strftime(time_format)

    return res


# 单次获取币本位合约历史K线数据
def single_fetch_coin_futures_candle_data(exchange, symbol, contractType, time_interval, start_time_utc, end_time_utc):
    """
    :param exchange:币安交易所
    :param symbol: 交易对，如：BTCUSD、ETHUSD
    :param contractType: 合约类型
    :param time_interval: K线时间间隔
    :param start_time_utc: K线的开始时间
    :param end_time_utc: K线的结束时间
    :return:
    """

    start_time_unix = int(pd.to_datetime(start_time_utc).timestamp() * 1000)
    end_time_unix = int(pd.to_datetime(end_time_utc).timestamp() * 1000)  # 转化为时间戳(ms)

    df_list = []
    while True:
        try:
            params = {
                'pair': symbol,
                'contractType': contractType,
                'interval': time_interval,
                'endTime': end_time_unix,
                'limit': 1000,  # 币安交易所规定默认500，最大1500
            }
            # 连续合约K线数据
            # GET /dapi/v1/continuousKlines
            df = pd.DataFrame(exchange.dapiPublicGetContinuousKlines(params=params))

            # 合并数据
            df_list.append(df)
            # 将本次查询K线结尾的时间，作为下一次查询K线开头的时间
            next_start_time_unix = int(df.iloc[0][0])
            end_time_unix = next_start_time_unix  # 将本周期最早K线时间作为下一次endTime时间
            _next_start_time_utc = time_strptime_timestamp(next_start_time_unix, '%Y-%m-%d %H:%M')

            next_end_time_unix = int(df.iloc[-1][0])
            _next_end_time_utc = time_strptime_timestamp(next_end_time_unix, '%Y-%m-%d %H:%M')
            print(f'      获取到 {_next_start_time_utc} -> {_next_end_time_utc} 的K线数据')
            if next_start_time_unix <= start_time_unix:  # 早去指定的K线最早开始时间，就结束循环
                break
            time.sleep(1)  # 抓取间隔需要暂停1s，防止抓取过于频繁
        except Exception as e:
            print('     报错：', str(e))
            time.sleep(1)
            break

    # ===合并整理数据
    df = pd.concat(df_list, ignore_index=True)
    df.rename(columns={0: 'MTS', 1: 'open', 2: 'high', 3: 'low', 4: 'close', 5: 'volume', 7: 'quote_volume',
                       8: 'trade_num', 9: 'taker_buy_base_asset_volume', 10: 'taker_buy_quote_as_set_volume'},
              inplace=True)  # 重命名
    df['MTS'] = df['MTS'].apply(lambda x: int(x))  # str类型，直接使用pd.to_datetime(df['MTS'], unit='ms') 会有精度偏差，故先要转换为int
    df['candle_begin_time'] = pd.to_datetime(df['MTS'], unit='ms')
    # 去除时间转换后遗留的小数，如：2022-05-28 07:40:27.520000
    df['candle_begin_time'] = df['candle_begin_time'].apply(lambda x: x.floor(freq='s'))
    df = df[df['candle_begin_time'] >= pd.to_datetime(start_time_utc)]
    df = df[['candle_begin_time', 'open', 'high', 'low', 'close', 'volume', 'quote_volume',
             'taker_buy_base_asset_volume', 'taker_buy_quote_as_set_volume', 'trade_num', ]]  # 整理列的顺序
    # df = df[df['candle_begin_time'].dt.date == pd.to_datetime(start_time_utc).date()]  # 选取数据时间段
    # 去重、排序
    df.drop_duplicates(subset=['candle_begin_time'], keep='last', inplace=True)
    df.sort_values('candle_begin_time', inplace=True)
    df.reset_index(drop=True, inplace=True)

    save_kline_data(df, symbol, contractType, time_interval)


# 将获取到的k先数据保存
def save_kline_data(df, symbol, contractType, time_interval):
    # 创建文件夹
    coin_futures_path = ''
    if contractType == 'PERPETUAL' and time_interval.endswith('m'):
        coin_futures_path = 'data/币安币本位永续合约1分钟k线数据'
    if contractType == 'PERPETUAL' and time_interval.endswith('h'):
        coin_futures_path = 'data/币安币本位永续合约1小时k线数据'
    elif contractType == 'CURRENT_QUARTER' and time_interval.endswith('m'):
        coin_futures_path = 'data/币安币本位交割合约1分钟k线数据'
    elif contractType == 'CURRENT_QUARTER' and time_interval.endswith('h'):
        coin_futures_path = 'data/币安币本位交割合约1小时k线数据'
    if os.path.exists(coin_futures_path) is False:  # path是文件夹或者文件的相对路径或者绝对路径
        os.makedirs(coin_futures_path)
    file_name = symbol + '.csv'
    # 拼接文件目录
    coin_futures_path = os.path.join(coin_futures_path, file_name)
    # 保存数据
    df.to_csv(coin_futures_path, index=False)


# 获取并保存币本位合约历史K线数据
def fetch_coin_futures_candle_data(exchange, symbol_list, contractType,
                                   time_interval_list, start_time_utc, end_time_utc):
    _time = datetime.now()  # 标记开始时间
    error_list = []  # 记录错误信息
    print(f'开始抓取{exchange.id}交易所：')
    for symbol in symbol_list:  # 遍历交易对
        for time_interval in time_interval_list:  # 遍历时间周期
            _time = datetime.now()  # 标记开始时间
            print(
                f'   {symbol}交易对 {start_time_utc} -> {end_time_utc}时间 {contractType}合约 {time_interval}时间周期 的历史K线数据')
            try:
                single_fetch_coin_futures_candle_data(
                    exchange, symbol, contractType, time_interval, start_time_utc, end_time_utc)
            except Exception as e:
                error_list.append('_'.join([symbol, contractType, time_interval, start_time_utc, end_time_utc]))
            finally:
                print(f'   当前周期获取完毕，共计花费：{datetime.now() - _time}\n')
    print(f'获取历史K线数据完毕，共计花费：{datetime.now() - _time}\n失败信息为：\n{error_list}')


if __name__ == '__main__':
    fetch_coin_futures_candle_data(exchange, symbol_list, contractType,
                                   time_interval_list, start_time_utc, end_time_utc)
